import os
import requests
from datetime import datetime
import base64
import subprocess
import time
import sys
from tqdm import tqdm
import threading

# 获取环境变量中的路径和 API Key
bozo = os.environ.get("BOZO") + "/Downloads/"
os.chdir(bozo)

def generate_image(prompt, api_key, image_url=None):
    """使用Gemini模型生成图像，支持纯文本提示或基于现有图像的修改"""
    start_time = time.time()
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    # 构建内容数组
    content = [{"type": "text", "text": prompt}]
    if image_url:
        content.append({
            "type": "image_url",
            "image_url": {"url": image_url}
        })
    
    data = {
        "model": "google/gemini-2.5-flash-image-preview:free",
        "temperature": 0.5,
        "messages": [{"role": "user", "content": content}]
    }
    
    try:
        print("正在生成图像，请稍候...")
        progress_bar = tqdm(total=100, bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt}')
        request_complete = False  # 用于跟踪请求是否完成
        
        # 修改：使用外部变量跟踪请求状态，而不是直接访问response
        def update_progress():
            while not request_complete:
                if progress_bar.n < 95:  # 留5%最后完成
                    progress_bar.update(5)
                time.sleep(0.3)
        
        # 启动进度条更新线程
        thread = threading.Thread(target=update_progress)
        thread.start()
        
        # 发送API请求
        response = requests.post(url, json=data, headers=headers)
        response.raise_for_status()
        result = response.json()
        
        # 请求完成，更新状态变量
        request_complete = True
        thread.join()  # 等待进度条线程结束
        progress_bar.update(100 - progress_bar.n)  # 确保进度条完成
        progress_bar.close()
        
        # 提取图像URL
        image_url = result["choices"][0]["message"]["images"][0]["image_url"]["url"]
        
        # 计算用时
        elapsed_time = time.time() - start_time
        print(f"图像生成完成，用时: {elapsed_time:.2f}秒")
        
        return image_url
    except Exception as e:
        # 确保进度条正确关闭
        request_complete = True  # 通知线程结束
        if 'thread' in locals():
            thread.join()
        if 'progress_bar' in locals():
            progress_bar.close()
        
        print(f"生成图像出错: {str(e)}")
        if 'response' in locals() and response:
            print(f"响应内容: {response.text}")
        
        elapsed_time = time.time() - start_time
        print(f"操作失败，用时: {elapsed_time:.2f}秒")
        return None

def save_image_from_data_url(data_url, save_path):
    """从data URL保存图像到本地文件"""
    start_time = time.time()
    
    try:
        # print("正在保存图像...")
        progress_bar = tqdm(total=100, bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt}')
        
        if data_url.startswith('data:image/'):
            header, data = data_url.split(',', 1)
            progress_bar.update(30)
            
            image_data = base64.b64decode(data)
            progress_bar.update(40)
            
            with open(save_path, 'wb') as f:
                f.write(image_data)
            progress_bar.update(30)
            
            progress_bar.close()
            print(f"图像已保存到: {save_path}")
            
            elapsed_time = time.time() - start_time
            # print(f"图像保存完成，用时: {elapsed_time:.2f}秒")
            
            return True
        else:
            progress_bar.close()
            print("无效的图像数据URL")
            return False
    except Exception as e:
        if 'progress_bar' in locals():
            progress_bar.close()
        print(f"保存图像出错: {str(e)}")
        
        elapsed_time = time.time() - start_time
        print(f"保存失败，用时: {elapsed_time:.2f}秒")
        return False

def open_image_with_os(save_path):
    """使用系统默认程序打开图像文件（跨平台兼容）"""
    try:
        if os.name == 'nt':
            os.startfile(save_path)
        elif os.name == 'posix':
            subprocess.run(['open' if sys.platform == 'darwin' else 'xdg-open', save_path])
    except Exception as e:
        print(f"打开图像失败: {str(e)}")

def main():
    total_start_time = time.time()
    
    deepseek_api_key = os.environ.get("MODELSCOPE_SDK_TOKEN")
    gemini_api_key = os.environ.get("ORAPI")
    
    prompt = input("请输入修改的英文提示词：\r\n")
    image_url = input("请输入图片URL：\r\n")
    
    image_url = generate_image(prompt, gemini_api_key, image_url)
    if not image_url:
        print("无法生成图像，程序退出")
        return
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"OR_{timestamp}.png"
    save_path = os.path.join(os.getcwd(), filename)
    if not save_image_from_data_url(image_url, save_path):
        print("无法保存图像，程序退出")
        return
    
    open_image_with_os(save_path)
    
    total_elapsed_time = time.time() - total_start_time
    # print(f"全部操作完成，总用时: {total_elapsed_time:.2f}秒")

if __name__ == "__main__":
    main()
    
